package controller;

import java.util.ArrayList;

import model.AIDataNode;


/**
 * Implementation of NaiveBayes Algorithm
 * @author walonka,zyner
 *
 */
public class NaiveBayes {
	private ArrayList<AIDataNode> trainingSet;
	private ArrayList<AIDataNode> testSet;
	private double[] class1std;
	private double[] class1mean;
	private double[] class0std;
	private double[] class0mean;
	
	/**
	 * Constructor of the Class
	 * @param trainingSet Set for the trainingData
	 * @param testSet set for the testingData
	 */
	public NaiveBayes(ArrayList<AIDataNode> trainingSet,ArrayList<AIDataNode> testSet){
		this.trainingSet=trainingSet;
		this.testSet=testSet;
		class1std = new double[8];
		class0std = new double[8];
		class1mean = new double[8];
		class0mean = new double[8];


	}

	/**
	 * Function to train the BayesClassifier using the Trainingdata
	 */
	public void train(){
		int num0=0;
		int num1=0;
		for(int i =0; i<trainingSet.size();i++){
			if(trainingSet.get(i).getKlass()==0){
				//Yeay, class0
				num0++;
				for (int j=0;j<trainingSet.get(i).getNodeWeight().length;j++){
					class0mean[j]+=trainingSet.get(i).getNodeWeight()[j];
				}
			}
			else{
				num1++;
				for (int j=0;j<trainingSet.get(i).getNodeWeight().length;j++){
					class1mean[j]+=trainingSet.get(i).getNodeWeight()[j];
				}
				
			}
		}		
		for(int i=0;i<8; i++){
			try{
				class0mean[i]/=num0;
			}
			catch(Exception e){
				
			}
			try{
				class1mean[i]/=num1;
			}
			catch(Exception e){
				
			}
		}

		for(int i =0; i<trainingSet.size();i++){
			if(trainingSet.get(i).getKlass()==0){
				//Yeay, class0
				//num0++;
				for (int j=0;j<trainingSet.get(i).getNodeWeight().length;j++){
					class0std[j]+=Math.pow((trainingSet.get(i).getNodeWeight()[j]-class0mean[j]),2);
				}


			}
			else{
				//num1++;
				for (int j=0;j<trainingSet.get(i).getNodeWeight().length;j++){
					class1std[j]+=Math.pow((trainingSet.get(i).getNodeWeight()[j]-class1mean[j]),2);
				}				
			}
		}	
		for(int i=0;i<8; i++){
			try{
				class0std[i]=Math.sqrt(class0std[i]/num0);
			}
			catch(Exception e){
				
			}
			try {
				class1std[i]=Math.sqrt(class1std[i]/num1);
			} catch (Exception e) {
			}

		}
		

	}

	/**
	 * Function to test the trained data
	 * @return returns the result of testing, first [0] the correct classification, second [1] the wrong classification.
	 */
	public int[] test(){
		int numCorrect = 0;
		int numWrong = 0;
		for (int i=0; i < testSet.size(); i++){
			double class0prob = 1;
			double class1prob = 1;
			for (int j = 0; j < testSet.get(i).getNodeWeight().length; j++){
				class0prob *= pdf(class0mean[j],class0std[j],testSet.get(i).getNodeWeight()[j]);
				class1prob *= pdf(class1mean[j],class1std[j],testSet.get(i).getNodeWeight()[j]);
			}
			if (class0prob>class1prob){ //node has been classified as class 0
				if(testSet.get(i).getKlass() == 0){
					numCorrect++;
				} else{
					numWrong++;
				}
			} else{
				if(testSet.get(i).getKlass() == 1){
					numCorrect++;
				} else{
					numWrong++;
				}
				
			}
		}
		int returndata[] = {numCorrect,numWrong};
		return returndata;
	}

	/**
	 * Function to calculate the probability using the formula of the lecture.
	 * @param mean
	 * @param std
	 * @param value
	 * @return probability
	 */
	public double pdf(double mean, double std, double value){
		return ((1/(std*Math.sqrt(2*Math.PI))) 
				* Math.exp(
						(-(Math.pow(value - mean, 2))/
								(2*(Math.pow(std, 2))))
							));			
	}
}

